import pygame
import random
import sys
import math
from enum import Enum

# 初始化pygame
pygame.init()
pygame.mixer.init()  # 初始化音效系统

# 游戏窗口尺寸
SCREEN_WIDTH = 600
SCREEN_HEIGHT = 800
screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
pygame.display.set_caption("Aircraft Battle")

# 颜色定义
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
RED = (255, 0, 0)
GREEN = (0, 255, 0)
BLUE = (0, 0, 255)
YELLOW = (255, 255, 0)
ORANGE = (255, 165, 0)
PURPLE = (128, 0, 128)
GRAY = (128, 128, 128)
LIGHT_BLUE = (173, 216, 230)

# 游戏设置
FPS = 60
ENEMY_SPAWN_RATE = 45  # 初始敌机生成速率
MAX_BULLETS = 15
DIFFICULTY_INCREASE_INTERVAL = 300  # 每5秒增加难度

# 游戏时钟
clock = pygame.time.Clock()

class GameState(Enum):
    MENU = 1
    PLAYING = 2
    PAUSED = 3
    GAME_OVER = 4
    SHOP = 5

# 武器类型
class WeaponType(Enum):
    BASIC = 1
    DOUBLE = 2
    TRIPLE = 3
    LASER = 4
    SPREAD = 5

# 升级类型
class UpgradeType(Enum):
    WEAPON = 1
    SPEED = 2
    HEALTH = 3
    FIRE_RATE = 4
    DAMAGE = 5

# 星星背景类
class Star:
    def __init__(self):
        self.x = random.randint(0, SCREEN_WIDTH)
        self.y = random.randint(0, SCREEN_HEIGHT)
        self.speed = random.uniform(0.5, 2)
        self.brightness = random.randint(100, 255)
        self.size = random.randint(1, 2)

    def update(self):
        self.y += self.speed
        if self.y > SCREEN_HEIGHT:
            self.y = -10
            self.x = random.randint(0, SCREEN_WIDTH)

    def draw(self):
        color = (self.brightness, self.brightness, self.brightness)
        pygame.draw.circle(screen, color, (int(self.x), int(self.y)), self.size)

# 爆炸效果类
class Explosion:
    def __init__(self, x, y, size="medium"):
        self.x = x
        self.y = y
        self.size = size
        self.radius = 5 if size == "small" else 10
        self.max_radius = 15 if size == "small" else 30
        self.growth_rate = 2 if size == "small" else 3
        self.alpha = 255
        self.particles = []

        # 创建粒子效果
        for _ in range(10 if size == "small" else 20):
            angle = random.uniform(0, 2 * math.pi)
            speed = random.uniform(1, 5)
            self.particles.append({
                'x': x, 'y': y,
                'vx': math.cos(angle) * speed,
                'vy': math.sin(angle) * speed,
                'life': 30,
                'color': random.choice([YELLOW, ORANGE, RED])
            })

    def update(self):
        self.radius += self.growth_rate
        self.alpha = max(0, 255 - (self.radius / self.max_radius) * 255)

        # 更新粒子
        for particle in self.particles[:]:
            particle['x'] += particle['vx']
            particle['y'] += particle['vy']
            particle['life'] -= 1
            if particle['life'] <= 0:
                self.particles.remove(particle)

        return self.radius < self.max_radius or len(self.particles) > 0

    def draw(self):
        # 绘制爆炸圆圈
        if self.alpha > 0:
            surf = pygame.Surface((self.max_radius * 2, self.max_radius * 2), pygame.SRCALPHA)
            color = (*ORANGE, int(self.alpha))
            pygame.draw.circle(surf, color, (self.max_radius, self.max_radius), int(self.radius))
            screen.blit(surf, (self.x - self.max_radius, self.y - self.max_radius))

        # 绘制粒子
        for particle in self.particles:
            alpha = int((particle['life'] / 30) * 255)
            pygame.draw.circle(screen, particle['color'],
                             (int(particle['x']), int(particle['y'])), 2)

# 升级商店类
class UpgradeShop:
    def __init__(self):
        self.upgrades = {
            UpgradeType.WEAPON: {"cost": 100, "level": 0, "max_level": 4, "name": "Weapon"},
            UpgradeType.SPEED: {"cost": 50, "level": 0, "max_level": 5, "name": "Speed"},
            UpgradeType.HEALTH: {"cost": 75, "level": 0, "max_level": 5, "name": "Max Health"},
            UpgradeType.FIRE_RATE: {"cost": 60, "level": 0, "max_level": 5, "name": "Fire Rate"},
            UpgradeType.DAMAGE: {"cost": 80, "level": 0, "max_level": 3, "name": "Damage"}
        }
        self.selected_index = 0

    def get_upgrade_cost(self, upgrade_type):
        upgrade = self.upgrades[upgrade_type]
        return int(upgrade["cost"] * (1.5 ** upgrade["level"]))

    def can_upgrade(self, upgrade_type, coins):
        upgrade = self.upgrades[upgrade_type]
        return upgrade["level"] < upgrade["max_level"] and coins >= self.get_upgrade_cost(upgrade_type)

    def upgrade(self, upgrade_type, coins):
        if self.can_upgrade(upgrade_type, coins):
            cost = self.get_upgrade_cost(upgrade_type)
            self.upgrades[upgrade_type]["level"] += 1
            return cost
        return 0

    def get_weapon_type(self, level):
        weapon_map = {0: WeaponType.BASIC, 1: WeaponType.DOUBLE, 2: WeaponType.TRIPLE, 3: WeaponType.LASER, 4: WeaponType.SPREAD}
        return weapon_map.get(level, WeaponType.BASIC)

# 玩家飞机类
class Player:
    def __init__(self):
        self.width = 50
        self.height = 60
        self.x = SCREEN_WIDTH // 2 - self.width // 2
        self.y = SCREEN_HEIGHT - self.height - 50
        self.base_speed = 8
        self.speed = self.base_speed
        self.base_max_health = 100
        self.max_health = self.base_max_health
        self.health = self.max_health
        self.base_fire_cooldown = 8
        self.fire_cooldown = 0
        self.base_damage = 1
        self.damage = self.base_damage
        self.invulnerable_timer = 0

        # 升级系统
        self.weapon_type = WeaponType.BASIC
        self.speed_level = 0
        self.health_level = 0
        self.fire_rate_level = 0
        self.damage_level = 0

        self.create_image()

    def create_image(self):
        self.image = pygame.Surface((self.width, self.height), pygame.SRCALPHA)
        # 根据武器类型改变飞机颜色
        if self.weapon_type == WeaponType.LASER:
            base_color = BLUE
            wing_color = (0, 0, 200)
        elif self.weapon_type == WeaponType.SPREAD:
            base_color = PURPLE
            wing_color = (100, 0, 100)
        elif self.weapon_type == WeaponType.TRIPLE:
            base_color = ORANGE
            wing_color = (200, 100, 0)
        elif self.weapon_type == WeaponType.DOUBLE:
            base_color = YELLOW
            wing_color = (200, 200, 0)
        else:
            base_color = GREEN
            wing_color = (0, 200, 0)

        # 机身
        pygame.draw.polygon(self.image, base_color,
                          [(25, 0), (5, 50), (45, 50)])
        # 机翼
        pygame.draw.polygon(self.image, wing_color,
                          [(0, 30), (25, 40), (50, 30)])
        # 座舱
        pygame.draw.ellipse(self.image, LIGHT_BLUE, (15, 10, 20, 15))
        # 引擎火焰
        pygame.draw.polygon(self.image, ORANGE,
                          [(15, 50), (20, 60), (25, 50)])
        pygame.draw.polygon(self.image, ORANGE,
                          [(25, 50), (30, 60), (35, 50)])

    def update(self):
        if self.fire_cooldown > 0:
            self.fire_cooldown -= 1
        if self.invulnerable_timer > 0:
            self.invulnerable_timer -= 1

    def draw(self):
        # 无敌时闪烁效果
        if self.invulnerable_timer > 0 and self.invulnerable_timer % 10 < 5:
            screen.blit(self.image, (self.x, self.y))
        else:
            screen.blit(self.image, (self.x, self.y))

        # 绘制血条
        bar_width = 40
        bar_height = 4
        bar_x = self.x + 5
        bar_y = self.y - 10

        # 血条背景
        pygame.draw.rect(screen, RED, (bar_x, bar_y, bar_width, bar_height))
        # 当前血量
        health_width = int(bar_width * (self.health / self.max_health))
        pygame.draw.rect(screen, GREEN, (bar_x, bar_y, health_width, bar_height))

    def move(self, dx, dy):
        new_x = self.x + dx * self.speed
        new_y = self.y + dy * self.speed

        # 边界检查
        self.x = max(0, min(new_x, SCREEN_WIDTH - self.width))
        self.y = max(0, min(new_y, SCREEN_HEIGHT - self.height))

    def fire(self):
        if self.fire_cooldown <= 0:
            self.fire_cooldown = self.base_fire_cooldown
            bullets = []

            if self.weapon_type == WeaponType.BASIC:
                bullets.append(Bullet(self.x + self.width // 2, self.y, damage=self.damage))
            elif self.weapon_type == WeaponType.DOUBLE:
                bullets.append(Bullet(self.x + 10, self.y, damage=self.damage))
                bullets.append(Bullet(self.x + self.width - 10, self.y, damage=self.damage))
            elif self.weapon_type == WeaponType.TRIPLE:
                bullets.append(Bullet(self.x + self.width // 2, self.y, damage=self.damage))
                bullets.append(Bullet(self.x + 5, self.y + 5, damage=self.damage))
                bullets.append(Bullet(self.x + self.width - 5, self.y + 5, damage=self.damage))
            elif self.weapon_type == WeaponType.LASER:
                # 激光武器 - 更快更强
                bullet = Bullet(self.x + self.width // 2, self.y, damage=self.damage * 2)
                bullet.speed = 20
                bullet.width = 8
                bullet.create_image()
                bullets.append(bullet)
            elif self.weapon_type == WeaponType.SPREAD:
                # 散射武器 - 5发子弹
                for i in range(5):
                    angle = (i - 2) * 15  # -30, -15, 0, 15, 30度
                    bullet = Bullet(self.x + self.width // 2, self.y, damage=self.damage)
                    bullet.spread_angle = angle
                    bullets.append(bullet)

            return bullets
        return []

    def upgrade_speed(self, level):
        self.speed_level = level
        self.speed = self.base_speed + level * 2

    def upgrade_health(self, level):
        self.health_level = level
        self.max_health = self.base_max_health + level * 25
        self.health = self.max_health  # 升级时回满血

    def upgrade_fire_rate(self, level):
        self.fire_rate_level = level
        self.base_fire_cooldown = max(3, 8 - level * 2)

    def upgrade_damage(self, level):
        self.damage_level = level
        self.damage = self.base_damage + level

    def upgrade_weapon(self, weapon_type):
        self.weapon_type = weapon_type
        self.create_image()  # 重新创建飞机图像

    def take_damage(self, damage):
        if self.invulnerable_timer <= 0:
            self.health -= damage
            self.invulnerable_timer = 60  # 1秒无敌时间
            return True
        return False

    def get_rect(self):
        return pygame.Rect(self.x, self.y, self.width, self.height)

# 子弹类
class Bullet:
    def __init__(self, x, y, is_enemy=False, damage=1):
        self.x = x
        self.y = y
        self.width = 4
        self.height = 12
        self.speed = 12 if not is_enemy else 6
        self.is_enemy = is_enemy
        self.damage = damage
        self.spread_angle = 0  # 散射角度
        self.create_image()

    def create_image(self):
        self.image = pygame.Surface((self.width, self.height), pygame.SRCALPHA)
        if self.is_enemy:
            # 敌机子弹
            pygame.draw.ellipse(self.image, PURPLE, (0, 0, self.width, self.height))
            pygame.draw.ellipse(self.image, RED, (1, 1, self.width-2, self.height-2))
        else:
            # 玩家子弹 - 根据伤害调整颜色
            if self.damage >= 3:
                color = RED
                core_color = ORANGE
            elif self.damage >= 2:
                color = ORANGE
                core_color = YELLOW
            else:
                color = YELLOW
                core_color = WHITE

            pygame.draw.ellipse(self.image, color, (0, 0, self.width, self.height))
            pygame.draw.ellipse(self.image, core_color, (1, 2, self.width-2, 4))

    def update(self):
        if self.is_enemy:
            self.y += self.speed
        else:
            self.y -= self.speed

        # 散射角度偏移
        if hasattr(self, 'spread_angle') and self.spread_angle != 0:
            self.x += math.sin(math.radians(self.spread_angle)) * 5

    def draw(self):
        screen.blit(self.image, (self.x, self.y))

    def get_rect(self):
        return pygame.Rect(self.x, self.y, self.width, self.height)

    def is_off_screen(self):
        return self.y < -self.height or self.y > SCREEN_HEIGHT

# 敌机类
class Enemy:
    def __init__(self, enemy_type="basic"):
        self.type = enemy_type
        self.setup_enemy()

    def setup_enemy(self):
        if self.type == "basic":
            self.width = 40
            self.height = 40
            self.speed = random.uniform(2, 4)
            self.health = 1
            self.score_value = 10
            self.fire_rate = 0  # 基础敌机不射击
        elif self.type == "fast":
            self.width = 35
            self.height = 35
            self.speed = random.uniform(4, 6)
            self.health = 1
            self.score_value = 15
            self.fire_rate = 0
        elif self.type == "heavy":
            self.width = 50
            self.height = 50
            self.speed = random.uniform(1, 2)
            self.health = 3
            self.score_value = 25
            self.fire_rate = 120  # 重型敌机射击

        self.x = random.randint(0, SCREEN_WIDTH - self.width)
        self.y = -self.height
        self.fire_cooldown = 0
        self.create_image()

    def create_image(self):
        self.image = pygame.Surface((self.width, self.height), pygame.SRCALPHA)

        if self.type == "basic":
            # 基础敌机
            pygame.draw.polygon(self.image, RED,
                              [(self.width//2, self.height), (0, 0), (self.width, 0)])
            pygame.draw.polygon(self.image, (150, 0, 0),
                              [(self.width//2, self.height-5), (5, 5), (self.width-5, 5)])
        elif self.type == "fast":
            # 快速敌机
            pygame.draw.polygon(self.image, PURPLE,
                              [(self.width//2, self.height), (0, 0), (self.width, 0)])
            pygame.draw.circle(self.image, (200, 0, 200),
                             (self.width//2, self.height//2), 8)
        elif self.type == "heavy":
            # 重型敌机
            pygame.draw.rect(self.image, (100, 0, 0),
                           (0, 0, self.width, self.height))
            pygame.draw.rect(self.image, RED,
                           (5, 5, self.width-10, self.height-10))
            pygame.draw.circle(self.image, YELLOW,
                             (self.width//2, self.height//2), 5)

    def update(self):
        self.y += self.speed

        # 左右摆动
        if self.type == "fast":
            self.x += math.sin(self.y * 0.05) * 2

        # 射击逻辑
        if self.fire_rate > 0:
            self.fire_cooldown = max(0, self.fire_cooldown - 1)

    def can_fire(self):
        return self.fire_rate > 0 and self.fire_cooldown <= 0

    def fire(self):
        if self.can_fire():
            self.fire_cooldown = self.fire_rate
            bullets = []

            if self.type == "heavy":
                # 重型敌机 - 三发散弹
                for i in range(3):
                    angle = (i - 1) * 10  # -10, 0, 10度
                    bullet = Bullet(self.x + self.width // 2, self.y + self.height, is_enemy=True)
                    bullet.spread_angle = angle
                    bullet.speed = 5
                    bullets.append(bullet)
            elif self.type == "fast":
                # 快速敌机 - 快速单发
                bullet = Bullet(self.x + self.width // 2, self.y + self.height, is_enemy=True)
                bullet.speed = 8
                bullets.append(bullet)
            else:
                # 基础敌机 - 标准单发
                bullets.append(Bullet(self.x + self.width // 2, self.y + self.height, is_enemy=True))

            return bullets
        return []

    def take_damage(self, damage):
        self.health -= damage
        return self.health <= 0

    def draw(self):
        screen.blit(self.image, (self.x, self.y))

        # 绘制血条（重型敌机）
        if self.type == "heavy" and self.health > 1:
            bar_width = 30
            bar_height = 3
            bar_x = self.x + 5
            bar_y = self.y - 8

            pygame.draw.rect(screen, RED, (bar_x, bar_y, bar_width, bar_height))
            health_width = int(bar_width * (self.health / 3))
            pygame.draw.rect(screen, GREEN, (bar_x, bar_y, health_width, bar_height))

    def get_rect(self):
        return pygame.Rect(self.x, self.y, self.width, self.height)

    def is_off_screen(self):
        return self.y > SCREEN_HEIGHT + self.height

# 道具类
class PowerUp:
    def __init__(self, x, y, power_type):
        self.x = x
        self.y = y
        self.width = 30
        self.height = 30
        self.type = power_type  # "health", "rapid_fire", "shield"
        self.speed = 2
        self.create_image()

    def create_image(self):
        self.image = pygame.Surface((self.width, self.height), pygame.SRCALPHA)

        if self.type == "health":
            # 生命值道具
            pygame.draw.circle(self.image, RED, (15, 15), 12)
            pygame.draw.rect(self.image, WHITE, (12, 8, 6, 14))
            pygame.draw.rect(self.image, WHITE, (8, 12, 14, 6))
        elif self.type == "rapid_fire":
            # 快速射击道具
            pygame.draw.circle(self.image, YELLOW, (15, 15), 12)
            for i in range(3):
                pygame.draw.ellipse(self.image, ORANGE,
                                  (5 + i*8, 10, 6, 10))
        elif self.type == "shield":
            # 护盾道具
            pygame.draw.circle(self.image, BLUE, (15, 15), 12)
            pygame.draw.circle(self.image, LIGHT_BLUE, (15, 15), 8, 2)

    def update(self):
        self.y += self.speed

    def draw(self):
        screen.blit(self.image, (self.x, self.y))

    def get_rect(self):
        return pygame.Rect(self.x, self.y, self.width, self.height)

    def is_off_screen(self):
        return self.y > SCREEN_HEIGHT

# 游戏类
class Game:
    def __init__(self):
        self.reset_game()
        self.font_large = pygame.font.SysFont(None, 48)
        self.font_medium = pygame.font.SysFont(None, 36)
        self.font_small = pygame.font.SysFont(None, 24)

        # 创建星星背景
        self.stars = [Star() for _ in range(100)]

        # 游戏状态
        self.state = GameState.MENU

        # 难度控制
        self.difficulty_timer = 0
        self.enemy_spawn_rate = ENEMY_SPAWN_RATE
        self.max_enemies = 5

        # 商店系统
        self.shop = UpgradeShop()
        self.coins = 0
        self.wave = 1
        self.enemies_killed_this_wave = 0
        self.enemies_per_wave = 10

    def reset_game(self):
        self.player = Player()
        self.player_bullets = []
        self.enemy_bullets = []
        self.enemies = []
        self.powerups = []
        self.explosions = []
        self.score = 0
        self.high_score = getattr(self, 'high_score', 0)
        self.enemy_spawn_counter = 0
        self.powerup_spawn_counter = 0
        self.difficulty_timer = 0
        self.enemy_spawn_rate = ENEMY_SPAWN_RATE
        self.max_enemies = 5
        self.coins = 0
        self.wave = 1
        self.enemies_killed_this_wave = 0

    def apply_upgrades(self):
        # 应用商店中的升级到玩家
        weapon_level = self.shop.upgrades[UpgradeType.WEAPON]["level"]
        speed_level = self.shop.upgrades[UpgradeType.SPEED]["level"]
        health_level = self.shop.upgrades[UpgradeType.HEALTH]["level"]
        fire_rate_level = self.shop.upgrades[UpgradeType.FIRE_RATE]["level"]
        damage_level = self.shop.upgrades[UpgradeType.DAMAGE]["level"]

        self.player.upgrade_weapon(self.shop.get_weapon_type(weapon_level))
        self.player.upgrade_speed(speed_level)
        self.player.upgrade_health(health_level)
        self.player.upgrade_fire_rate(fire_rate_level)
        self.player.upgrade_damage(damage_level)

    def handle_events(self):
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                return False

            if event.type == pygame.KEYDOWN:
                if self.state == GameState.MENU:
                    if event.key == pygame.K_SPACE:
                        self.state = GameState.PLAYING
                        self.reset_game()
                    elif event.key == pygame.K_s:
                        self.state = GameState.SHOP
                    elif event.key == pygame.K_ESCAPE:
                        return False

                elif self.state == GameState.PLAYING:
                    if event.key == pygame.K_ESCAPE:
                        self.state = GameState.PAUSED
                    elif event.key == pygame.K_s:
                        self.state = GameState.SHOP

                elif self.state == GameState.PAUSED:
                    if event.key == pygame.K_ESCAPE:
                        self.state = GameState.PLAYING
                    elif event.key == pygame.K_q:
                        self.state = GameState.MENU
                    elif event.key == pygame.K_s:
                        self.state = GameState.SHOP

                elif self.state == GameState.SHOP:
                    if event.key == pygame.K_ESCAPE:
                        self.state = GameState.PLAYING
                    elif event.key == pygame.K_UP:
                        self.shop.selected_index = max(0, self.shop.selected_index - 1)
                    elif event.key == pygame.K_DOWN:
                        self.shop.selected_index = min(len(self.shop.upgrades) - 1, self.shop.selected_index + 1)
                    elif event.key == pygame.K_RETURN or event.key == pygame.K_SPACE:
                        self.purchase_upgrade()

                elif self.state == GameState.GAME_OVER:
                    if event.key == pygame.K_r:
                        self.state = GameState.PLAYING
                        self.reset_game()
                    elif event.key == pygame.K_ESCAPE:
                        self.state = GameState.MENU

        return True

    def purchase_upgrade(self):
        upgrade_types = list(self.shop.upgrades.keys())
        if self.shop.selected_index < len(upgrade_types):
            upgrade_type = upgrade_types[self.shop.selected_index]
            cost = self.shop.upgrade(upgrade_type, self.coins)
            if cost > 0:
                self.coins -= cost
                self.apply_upgrades()

    def update(self):
        if self.state != GameState.PLAYING:
            return

        # 更新星星背景
        for star in self.stars:
            star.update()

        # 更新玩家
        self.player.update()

        # 处理输入
        keys = pygame.key.get_pressed()
        dx = dy = 0
        if keys[pygame.K_LEFT] or keys[pygame.K_a]:
            dx = -1
        if keys[pygame.K_RIGHT] or keys[pygame.K_d]:
            dx = 1
        if keys[pygame.K_UP] or keys[pygame.K_w]:
            dy = -1
        if keys[pygame.K_DOWN]:  # 移除S键作为向下移动，避免与商店快捷键冲突
            dy = 1

        if dx != 0 or dy != 0:
            self.player.move(dx, dy)

        if keys[pygame.K_SPACE]:
            bullets = self.player.fire()
            for bullet in bullets:
                if len(self.player_bullets) < MAX_BULLETS:
                    self.player_bullets.append(bullet)

        # 更新子弹
        for bullet in self.player_bullets[:]:
            bullet.update()
            if bullet.is_off_screen():
                self.player_bullets.remove(bullet)

        for bullet in self.enemy_bullets[:]:
            bullet.update()
            if bullet.is_off_screen():
                self.enemy_bullets.remove(bullet)

        # 更新敌机
        for enemy in self.enemies[:]:
            enemy.update()

            # 敌机射击
            enemy_bullets = enemy.fire()
            self.enemy_bullets.extend(enemy_bullets)

            if enemy.is_off_screen():
                self.enemies.remove(enemy)

        # 更新道具
        for powerup in self.powerups[:]:
            powerup.update()
            if powerup.is_off_screen():
                self.powerups.remove(powerup)

        # 更新爆炸效果
        for explosion in self.explosions[:]:
            if not explosion.update():
                self.explosions.remove(explosion)

        # 碰撞检测 - 玩家子弹与敌机
        for bullet in self.player_bullets[:]:
            for enemy in self.enemies[:]:
                if bullet.get_rect().colliderect(enemy.get_rect()):
                    if bullet in self.player_bullets:
                        self.player_bullets.remove(bullet)

                    if enemy.take_damage(bullet.damage):
                        self.enemies.remove(enemy)
                        self.score += enemy.score_value
                        self.coins += enemy.score_value // 2  # 金币奖励
                        self.enemies_killed_this_wave += 1

                        self.explosions.append(Explosion(
                            enemy.x + enemy.width // 2,
                            enemy.y + enemy.height // 2,
                            "small"
                        ))

                        # 随机掉落道具
                        if random.random() < 0.15:  # 15%概率
                            power_type = random.choice(["health", "rapid_fire", "shield"])
                            self.powerups.append(PowerUp(
                                enemy.x + enemy.width // 2,
                                enemy.y + enemy.height // 2,
                                power_type
                            ))
                    break

        # 碰撞检测 - 敌机子弹与玩家
        for bullet in self.enemy_bullets[:]:
            if bullet.get_rect().colliderect(self.player.get_rect()):
                self.enemy_bullets.remove(bullet)
                if self.player.take_damage(10):
                    self.explosions.append(Explosion(
                        self.player.x + self.player.width // 2,
                        self.player.y + self.player.height // 2,
                        "small"
                    ))

        # 碰撞检测 - 敌机与玩家
        for enemy in self.enemies[:]:
            if enemy.get_rect().colliderect(self.player.get_rect()):
                self.enemies.remove(enemy)
                if self.player.take_damage(20):
                    self.explosions.append(Explosion(
                        self.player.x + self.player.width // 2,
                        self.player.y + self.player.height // 2,
                        "medium"
                    ))
                self.explosions.append(Explosion(
                    enemy.x + enemy.width // 2,
                    enemy.y + enemy.height // 2,
                    "medium"
                ))

        # 碰撞检测 - 道具与玩家
        for powerup in self.powerups[:]:
            if powerup.get_rect().colliderect(self.player.get_rect()):
                self.powerups.remove(powerup)
                self.apply_powerup(powerup.type)

        # 生成敌机
        self.enemy_spawn_counter += 1
        if self.enemy_spawn_counter >= self.enemy_spawn_rate:
            self.spawn_enemy()
            self.enemy_spawn_counter = 0

        # 增加难度
        self.difficulty_timer += 1
        if self.difficulty_timer >= DIFFICULTY_INCREASE_INTERVAL:
            self.increase_difficulty()
            self.difficulty_timer = 0

        # 检查波次完成
        if self.enemies_killed_this_wave >= self.enemies_per_wave:
            self.wave += 1
            self.enemies_killed_this_wave = 0
            self.enemies_per_wave += 2
            # 波次奖励
            wave_bonus = self.wave * 50
            self.coins += wave_bonus
            self.score += wave_bonus

        # 检查游戏结束
        if self.player.health <= 0:
            self.state = GameState.GAME_OVER
            if self.score > self.high_score:
                self.high_score = self.score
            self.explosions.append(Explosion(
                self.player.x + self.player.width // 2,
                self.player.y + self.player.height // 2,
                "large"
            ))

    def spawn_enemy(self):
        if len(self.enemies) >= self.max_enemies:
            return

        # 根据难度选择敌机类型
        rand = random.random()
        if self.difficulty_timer < 600:  # 前10秒只有基础敌机
            enemy_type = "basic"
        elif self.difficulty_timer < 1200:  # 10-20秒加入快速敌机
            enemy_type = "basic" if rand < 0.7 else "fast"
        else:  # 20秒后加入重型敌机
            if rand < 0.5:
                enemy_type = "basic"
            elif rand < 0.8:
                enemy_type = "fast"
            else:
                enemy_type = "heavy"

        self.enemies.append(Enemy(enemy_type))

    def increase_difficulty(self):
        self.enemy_spawn_rate = max(20, self.enemy_spawn_rate - 2)
        self.max_enemies = min(10, self.max_enemies + 1)

    def apply_powerup(self, power_type):
        if power_type == "health":
            self.player.health = min(self.player.max_health, self.player.health + 30)
        elif power_type == "rapid_fire":
            self.player.fire_cooldown = max(2, self.player.fire_cooldown - 2)
        elif power_type == "shield":
            self.player.invulnerable_timer = 180  # 3秒无敌

    def draw_menu(self):
        screen.fill(BLACK)

        # 绘制星星背景
        for star in self.stars:
            star.draw()

        # 标题
        title = self.font_large.render("AIRCRAFT BATTLE", True, WHITE)
        title_rect = title.get_rect(center=(SCREEN_WIDTH // 2, 150))
        screen.blit(title, title_rect)

        # 菜单选项
        start_text = self.font_medium.render("Press SPACE to Start", True, GREEN)
        start_rect = start_text.get_rect(center=(SCREEN_WIDTH // 2, 300))
        screen.blit(start_text, start_rect)

        shop_text = self.font_medium.render("Press S for Shop", True, YELLOW)
        shop_rect = shop_text.get_rect(center=(SCREEN_WIDTH // 2, 350))
        screen.blit(shop_text, shop_rect)

        quit_text = self.font_medium.render("Press ESC to Quit", True, RED)
        quit_rect = quit_text.get_rect(center=(SCREEN_WIDTH // 2, 400))
        screen.blit(quit_text, quit_rect)

        # 操作说明
        controls = [
            "Controls:",
            "Arrow Keys/WAD - Move",
            "SPACE - Fire",
            "S - Shop",
            "ESC - Pause"
        ]

        y = 450
        for line in controls:
            text = self.font_small.render(line, True, WHITE)
            text_rect = text.get_rect(center=(SCREEN_WIDTH // 2, y))
            screen.blit(text, text_rect)
            y += 30

    def draw_game(self):
        screen.fill(BLACK)

        # 绘制星星背景
        for star in self.stars:
            star.draw()

        # 绘制游戏对象
        self.player.draw()

        for bullet in self.player_bullets:
            bullet.draw()

        for bullet in self.enemy_bullets:
            bullet.draw()

        for enemy in self.enemies:
            enemy.draw()

        for powerup in self.powerups:
            powerup.draw()

        for explosion in self.explosions:
            explosion.draw()

        # 绘制UI
        score_text = self.font_medium.render(f"Score: {self.score}", True, WHITE)
        screen.blit(score_text, (10, 10))

        coins_text = self.font_small.render(f"Coins: {self.coins}", True, YELLOW)
        screen.blit(coins_text, (10, 50))

        wave_text = self.font_small.render(f"Wave: {self.wave}", True, ORANGE)
        screen.blit(wave_text, (10, 80))

        enemies_text = self.font_small.render(f"Enemies: {self.enemies_killed_this_wave}/{self.enemies_per_wave}", True, WHITE)
        screen.blit(enemies_text, (10, 110))

        # 显示武器类型
        weapon_names = {
            WeaponType.BASIC: "Basic",
            WeaponType.DOUBLE: "Double",
            WeaponType.TRIPLE: "Triple",
            WeaponType.LASER: "Laser",
            WeaponType.SPREAD: "Spread"
        }
        weapon_text = self.font_small.render(f"Weapon: {weapon_names[self.player.weapon_type]}", True, GREEN)
        screen.blit(weapon_text, (10, 140))

        # 显示难度等级
        difficulty_level = self.difficulty_timer // DIFFICULTY_INCREASE_INTERVAL + 1
        diff_text = self.font_small.render(f"Level: {difficulty_level}", True, RED)
        screen.blit(diff_text, (10, 170))

        # 商店提示
        shop_hint = self.font_small.render("Press S - Shop", True, GRAY)
        screen.blit(shop_hint, (SCREEN_WIDTH - 120, 10))

    def draw_paused(self):
        self.draw_game()

        # 半透明遮罩
        overlay = pygame.Surface((SCREEN_WIDTH, SCREEN_HEIGHT))
        overlay.set_alpha(128)
        overlay.fill(BLACK)
        screen.blit(overlay, (0, 0))

        # 暂停文字
        paused_text = self.font_large.render("PAUSED", True, WHITE)
        paused_rect = paused_text.get_rect(center=(SCREEN_WIDTH // 2, SCREEN_HEIGHT // 2 - 50))
        screen.blit(paused_text, paused_rect)

        resume_text = self.font_medium.render("Press ESC to Resume", True, GREEN)
        resume_rect = resume_text.get_rect(center=(SCREEN_WIDTH // 2, SCREEN_HEIGHT // 2 + 20))
        screen.blit(resume_text, resume_rect)

        shop_text = self.font_medium.render("Press S for Shop", True, YELLOW)
        shop_rect = shop_text.get_rect(center=(SCREEN_WIDTH // 2, SCREEN_HEIGHT // 2 + 60))
        screen.blit(shop_text, shop_rect)

        quit_text = self.font_medium.render("Press Q to Quit to Menu", True, RED)
        quit_rect = quit_text.get_rect(center=(SCREEN_WIDTH // 2, SCREEN_HEIGHT // 2 + 100))
        screen.blit(quit_text, quit_rect)

    def draw_game_over(self):
        self.draw_game()

        # 半透明遮罩
        overlay = pygame.Surface((SCREEN_WIDTH, SCREEN_HEIGHT))
        overlay.set_alpha(128)
        overlay.fill(BLACK)
        screen.blit(overlay, (0, 0))

        # 游戏结束文字
        game_over_text = self.font_large.render("GAME OVER", True, RED)
        game_over_rect = game_over_text.get_rect(center=(SCREEN_WIDTH // 2, SCREEN_HEIGHT // 2 - 100))
        screen.blit(game_over_text, game_over_rect)

        score_text = self.font_medium.render(f"Final Score: {self.score}", True, WHITE)
        score_rect = score_text.get_rect(center=(SCREEN_WIDTH // 2, SCREEN_HEIGHT // 2 - 30))
        screen.blit(score_text, score_rect)

        if self.score == self.high_score:
            high_score_text = self.font_medium.render("NEW HIGH SCORE!", True, YELLOW)
            high_score_rect = high_score_text.get_rect(center=(SCREEN_WIDTH // 2, SCREEN_HEIGHT // 2 + 10))
            screen.blit(high_score_text, high_score_rect)

        restart_text = self.font_medium.render("Press R to Restart", True, GREEN)
        restart_rect = restart_text.get_rect(center=(SCREEN_WIDTH // 2, SCREEN_HEIGHT // 2 + 60))
        screen.blit(restart_text, restart_rect)

        menu_text = self.font_medium.render("Press ESC for Menu", True, WHITE)
        menu_rect = menu_text.get_rect(center=(SCREEN_WIDTH // 2, SCREEN_HEIGHT // 2 + 100))
        screen.blit(menu_text, menu_rect)

    def draw_shop(self):
        screen.fill(BLACK)

        # 绘制星星背景
        for star in self.stars:
            star.draw()

        # 标题
        title = self.font_large.render("UPGRADE SHOP", True, YELLOW)
        title_rect = title.get_rect(center=(SCREEN_WIDTH // 2, 50))
        screen.blit(title, title_rect)

        # 显示金币
        coins_text = self.font_medium.render(f"Coins: {self.coins}", True, YELLOW)
        coins_rect = coins_text.get_rect(center=(SCREEN_WIDTH // 2, 100))
        screen.blit(coins_text, coins_rect)

        # 升级选项
        y = 180
        upgrade_types = list(self.shop.upgrades.keys())

        for i, upgrade_type in enumerate(upgrade_types):
            upgrade = self.shop.upgrades[upgrade_type]
            cost = self.shop.get_upgrade_cost(upgrade_type)
            can_afford = self.coins >= cost and upgrade["level"] < upgrade["max_level"]

            # 选中高亮
            if i == self.shop.selected_index:
                pygame.draw.rect(screen, WHITE, (140, y - 5, 320, 40), 2)

            # 升级名称和等级
            name_text = f"{upgrade['name']}: Lv.{upgrade['level']}/{upgrade['max_level']}"
            color = GREEN if can_afford else RED
            text = self.font_medium.render(name_text, True, color)
            screen.blit(text, (150, y))

            # 价格
            if upgrade["level"] < upgrade["max_level"]:
                cost_text = f"Cost: {cost}"
                cost_color = YELLOW if can_afford else RED
                cost_render = self.font_small.render(cost_text, True, cost_color)
                screen.blit(cost_render, (150, y + 30))
            else:
                max_text = self.font_small.render("MAXED", True, GREEN)
                screen.blit(max_text, (150, y + 30))

            # 武器特殊显示
            if upgrade_type == UpgradeType.WEAPON:
                weapon_names = ["Basic", "Double", "Triple", "Laser", "Spread"]
                if upgrade["level"] < len(weapon_names):
                    weapon_text = f"({weapon_names[upgrade['level']]})"
                    weapon_render = self.font_small.render(weapon_text, True, LIGHT_BLUE)
                    screen.blit(weapon_render, (350, y + 5))

            y += 70

        # 操作说明
        controls = [
            "UP/DOWN - Select",
            "ENTER/SPACE - Buy",
            "ESC - Back to Game"
        ]

        y = 550
        for line in controls:
            text = self.font_small.render(line, True, WHITE)
            text_rect = text.get_rect(center=(SCREEN_WIDTH // 2, y))
            screen.blit(text, text_rect)
            y += 30

    def draw(self):
        if self.state == GameState.MENU:
            self.draw_menu()
        elif self.state == GameState.PLAYING:
            self.draw_game()
        elif self.state == GameState.PAUSED:
            self.draw_paused()
        elif self.state == GameState.GAME_OVER:
            self.draw_game_over()
        elif self.state == GameState.SHOP:
            self.draw_shop()

        pygame.display.flip()

    def run(self):
        running = True
        while running:
            running = self.handle_events()
            self.update()
            self.draw()
            clock.tick(FPS)

        pygame.quit()
        sys.exit()

# 运行游戏
if __name__ == "__main__":
    game = Game()
    game.run()